library(tidyverse)
library(wesanderson)
library(MetBrewer)
library(sp)
library(sf)
library(classInt)
library(raster)
#library(exactextractr)
library(fixest)
library(rgdal)
library(raster)
library(classInt)
library(tidycensus)


add.alpha <- function(col, alpha=1){
  if(missing(col))
    stop("Please provide a vector of colours.")
  apply(sapply(col, col2rgb)/255, 2, 
        function(x) 
          rgb(x[1], x[2], x[3], alpha=alpha))  
}	
get_grid_raster <- function(raster){
  output <- raster(raster)
  output[] <- 1:ncell(output)
  return(output)
}
identify_non_cells <- function(raster, poly){
  raster = raster(raster)
  raster[] <- 1:ncell(raster)
  r_ext <- extract(raster, poly)
  all_cells <- 1:ncell(raster)
  non_cells <- all_cells[all_cells %in% unlist(r_ext) == F]
  return(non_cells)
  
}



load("figdat/fig_SI10_data.RData")

########################### plot ######################
  
  
  pal_pm <- colorRampPalette(colors = MetBrewer::met.brewer("VanGogh2", 8))(2000) %>% rev()
  pal_elev <- colorRampPalette(colors = MetBrewer::met.brewer("Tiepolo", 8))(2000) %>% rev()
  pal_wealth <- colorRampPalette(colors = c("red3",MetBrewer::met.brewer("Benedictus", 16)[c(1:6, 8:13)], "navy"))(2000) 
  
  
  
  
    pdf("figures/figure_SI10_maps.pdf", width = 12, height = 12)                
      
        par(mfrow = c(3,3))
        par(mar = c(4,0,2,0))
        
        
    
        
        ### la ###
    
        #[1]    
        plot(la_sp, col = NA, border = NA)
        plot(r_la_plot, col = pal_pm, axes = F,add=T, legend = F)
        plot(la_shp, add = T, lwd =2)
        mtext(side = 3, adj = 0, text = "PM2.5",cex=2, line = -2)
  
      
      
     
        
        #[2]
        int <- classIntervals(inc_la_sf2$estimate, style = "fixed", fixedBreaks = c(-10,seq(0, 75000, 1e4/5),1e6))
        col_la_w <- findColours(int, (pal_wealth))
        col_la_w[is.na(col_la_w)]<-pal_wealth[1000]
        
        plot(la_sp, border = NA)
        plot(st_geometry(inc_la_sf2), col = col_la_w, border = NA, add = T)
        plot(la_sp, add = T, lwd =2, bg = 'white')
        mtext(side = 3, adj = 0, text = "Wealth",cex=2, line = -2)
        
        
        
        #[3]
        
        
        plot(la_sp, col = NA, border = NA)
        plot(el_la, col = (pal_elev), axes = F,bty = "n",add=T, legend = F)
        plot(la_shp, add = T, lwd =2)
        mtext(side = 3, adj = 0, text = "Elevation",cex=2, line = -2)
        
        
        
        ###### Houston ########
        hou <- hou_sp
        #[1]
        plot(hou_sp, col = NA, border = NA)
        plot(r_hou_plot, col = pal_pm, axes = F,bty = "n",add=T, legend = F)
        plot(hou, add = T, lwd =1)
        mtext(side = 3, adj = 0, text = "PM2.5",cex=2, line = -2)
        
      
        #[2]
        
        int <- classIntervals(inc_hou_sf2$estimate, style = "fixed", fixedBreaks = c(-10,seq(0, 75000, 1e4/5),1e6))
        col_hou_w <- findColours(int, (pal_wealth))
        
        
        plot(st_geometry(inc_hou_sf2), col = col_hou_w, border = NA)
        plot(hou_shp, add = T, lwd =1)
        mtext(side = 3, adj = 0, text = "Wealth",cex=2, line = -2)
        
        
        #[3]
        
        plot(hou_sp, col = NA, border = NA)
        plot(el_hou, col = (pal_elev), axes = F,bty = "n",add=T, legend = F)
        plot(hou, add = T, lwd =1)
        mtext(side = 3, adj = 0, text = "Elevation",cex=2, line = -2)
         
  
        
        ####### NY #########  
        
        
      #[1]  
        plot(ny_sp, col = NA, border = NA)
        plot(r_ny_plot, col = pal_pm, axes = F,bty = "n",add=T, legend = F)
        plot(ny, add = T, lwd =2)
      
          mtext(side = 3, adj = 0, text = "PM2.5",cex=2, line = -2)

        
        
        plotrix::gradient.rect(-74.3,40.45,-73.7,40.475,
                               col=(plotrix::smoothColors(colorRampPalette(pal_pm)(2000))), 
                               gradient="x", border = 'black')
        
        
        
        #[2]
        
        int <- classIntervals(inc_ny$estimate, style = "fixed", fixedBreaks = c(-10,seq(0, 75000, 1e4/5),1e6))
        col_ny_w <- findColours(int, (pal_wealth))
        
        
        plot(st_geometry(inc_ny), col = col_ny_w, border = NA)
        plot(ny_sp, add = T, lwd =2)
        mtext(side = 3, adj = 0, text = "Wealth",cex=2, line = -2)
        
        
        
        
        plotrix::gradient.rect(-74.3,40.45,-73.7,40.475,
                               col=(plotrix::smoothColors(colorRampPalette(pal_wealth)(2000))), 
                               gradient="x", border = 'black')
        
        
        #[3]
        
        
        
        
        plot(ny_sp, col = NA, border = NA)
        plot(el_ny, col = (pal_elev), axes = F,bty = "n",add=T, legend = F)
        plot(ny, add = T, lwd =2)
        mtext(side = 3, adj = 0, text = "Elevation",cex=2, line = -2)
        
        
        
        plotrix::gradient.rect(-74.3,40.45,-73.7,40.475,
                               col=(plotrix::smoothColors(colorRampPalette(pal_elev)(2000))), 
                               gradient="x", border = 'black')
        
        
        
        dev.off()
        
        
        
        
        
        
        
        
        
        
        
        
        
        
        
        
  ########## scatters ###########
        
        
        
        #PM vs wealth
        
        #la
        inc_la2 <- inc_la_sf2 %>% drop_na(geometry)
        rloc_la <- data.frame(coordinates(as_Spatial(inc_la2)), wealth = inc_la2$estimate)
        rloc_la$pm <- r_la_plot[raster::cellFromXY(r_la_plot,xy = rloc_la[,1:2])]
        rloc_la$elev <- el_la[raster::cellFromXY(el_la,xy = rloc_la[,1:2])]
        names(rloc_la)<-c("longitude","latitude","wealth","pm","elev")
       
        inc_hou2 <- inc_hou_sf2 %>% drop_na(geometry)
        rloc_hou <- data.frame(coordinates(as_Spatial(inc_hou2)), wealth = inc_hou2$estimate)
        rloc_hou$pm <- r_hou_plot[raster::cellFromXY(r_hou_plot,xy = rloc_hou[,1:2])]
        rloc_hou$elev <- el_hou[raster::cellFromXY(el_hou,xy = rloc_hou[,1:2])]
        
        pmloc_hou = rloc_hou
        names(pmloc_hou)<-c("longitude","houtitude","wealth","pm","elev")
        pmloc_hou <- pmloc_hou %>% drop_na(pm, wealth, elev)
        scatter_pm_wealth_hou <- pmloc_hou
        
        inc_ny2 <- inc_ny %>% drop_na(geometry)
        rloc_ny <- data.frame(coordinates(as_Spatial(inc_ny2[!st_is_empty(inc_ny2),])), wealth = inc_ny2$estimate[!st_is_empty(inc_ny2)])
        rloc_ny$pm <- r_ny_plot[raster::cellFromXY(r_ny_plot,xy = rloc_ny[,1:2])]
        rloc_ny$elev <- el_ny[raster::cellFromXY(el_ny,xy = rloc_ny[,1:2])]
        
        pmloc_ny = rloc_ny
        names(pmloc_ny)<-c("longitude","nytitude","wealth","pm","elev")
        pmloc_ny <- pmloc_ny %>% drop_na(pm, wealth, elev)
        scatter_pm_wealth_ny <- pmloc_ny
        
        
      
      
  pdf("figures/figure_SI10_scatters.pdf", width = 12, height = 8)                
      
      par(mfrow =c(2,3))
      
      plot(rloc_la[,c("wealth","pm")], ylim = c(5,20), xlim = c(2500, 107000), axes = F, xlab = "", ylab = "", col = NA)
      points(rloc_la[,c("wealth","pm")], pch = 21, bg = add.alpha('gray70', .9), col = add.alpha('white', 1), cex=2, lwd = 0.25)
      lines(x  =seq(0,150000,1000), y =  predict(loess(pm ~ wealth, data = rloc_la), newdata = data.frame(wealth =  seq(0,150000,1000))), col  = 'blue',lwd=3)
      mtext(side = 2, text = "PM2.5",line=3,cex=2)
      mtext(side = 1, text = "Wealth",line=3,cex=2)
  
      
      plot(scatter_pm_wealth_hou[,c("wealth","pm")], axes = F, xlab = "", ylab = "", col = NA)
      points(scatter_pm_wealth_hou[,c("wealth","pm")], pch = 21, bg = add.alpha('gray70', .9), col = add.alpha('white', 1), cex=2, lwd = 0.25)
      lines(x  =seq(0,150000,1000), y =  predict(loess(pm ~ wealth, data = scatter_pm_wealth_hou), newdata = data.frame(wealth = seq(0,150000,1000))), col  = 'blue',lwd=3)
      mtext(side = 2, text = "PM2.5",line=3,cex=2)
      mtext(side = 1, text = "Wealth",line=3,cex=2)

      plot(scatter_pm_wealth_ny[,c("wealth","pm")], axes = F, xlab = "", ylab = "", col = NA)
      points(scatter_pm_wealth_ny[,c("wealth","pm")], pch = 21, bg = add.alpha('gray70', .9), col = add.alpha('white', 1), cex=2, lwd = 0.25)
      lines(x  =seq(0,150000,1000), y =  predict(loess(pm ~ wealth, data = scatter_pm_wealth_ny), newdata = data.frame(wealth =  seq(0,150000,1000))), col  = 'blue',lwd=3)
      mtext(side = 2, text = "PM2.5",line=3,cex=2)
      mtext(side = 1, text = "Wealth",line=3,cex=2)
     
       plot(rloc_la[,c("wealth","elev")],  axes = F, xlab = "", ylab = "", col = NA)
      points(rloc_la[,c("wealth","elev")], pch = 21, bg = add.alpha('gray70', .9), col = add.alpha('white', 1), cex=2, lwd = 0.25)
      lines(x =seq(0,150000,1000), y = predict(loess(elev ~ wealth, data = rloc_la), newdata= data.frame(wealth = seq(0,150000,1000))), col  = 'blue',lwd=3)
      mtext(side = 1, text = "Elevation",line=3,cex=2)
      mtext(side = 2, text = "Wealth",line=3,cex=2)
      
      
      plot(rloc_hou[,c("wealth","elev")], axes = F, xlab = "", ylab = "", col = NA)
      points(rloc_hou[,c("wealth","elev")], pch = 21, bg = add.alpha('gray70', .9), col = add.alpha('white', 1), cex=2, lwd = 0.25)
      lines(x = seq(0,150000,1000), y = predict(loess(elev ~ wealth, data = rloc_hou), newdata= data.frame(wealth = seq(0,150000,1000))), col  = 'blue',lwd=3)
      mtext(side = 1, text = "Wealth",line=3,cex=2)
      mtext(side = 2, text = "Elevation",line=3,cex=2)
      
      
      plot(rloc_ny[,c("wealth","elev")], axes = F, xlab = "", ylab = "", col = NA)
      points(rloc_ny[,c("wealth","elev")], pch = 21, bg = add.alpha('gray70', .9), col = add.alpha('white', 1), cex=2, lwd = 0.25)
      lines(x = seq(0,150000,1000), y = predict(loess(elev ~ wealth, data = rloc_ny), newdata= data.frame(wealth = seq(0,150000,1000))), col  = 'blue',lwd=3)
      mtext(side = 1, text = "Wealth",line=3,cex=2)
      mtext(side = 2, text = "Elevation",line=3,cex=2)
      
      
      
      dev.off()
      
        
      
    
        
        
